import json
import gradio as gr
import torch
import os

import yaml


#read GPU info
ngpu=torch.cuda.device_count()
gpu_infos=[]
if(torch.cuda.is_available() is False or ngpu==0):
    if_gpu_ok=False
else:
    if_gpu_ok = False
    for i in range(ngpu):
        gpu_name=torch.cuda.get_device_name(i)
        if("MX"in gpu_name):
            continue
        if("RTX" in gpu_name.upper() or "10"in gpu_name or "16"in gpu_name or "20"in gpu_name or "30"in gpu_name or "40"in gpu_name or "A50"in gpu_name.upper() or "70"in gpu_name or "80"in gpu_name or "90"in gpu_name or "M4"in gpu_name or"P4"in gpu_name or "T4"in gpu_name or "TITAN"in gpu_name.upper()):#A10#A100#V100#A40#P40#M40#K80
            if_gpu_ok=True#至少有一张能用的N卡
            gpu_infos.append("%s\t%s"%(i,gpu_name))
gpu_info="\n".join(gpu_infos)if if_gpu_ok is True and len(gpu_infos)>0 else "很遗憾您这没有能用的显卡来支持您训练"
gpus="-".join([i[0]for i in gpu_infos])

#read cuda info for inference
cuda = {}
if torch.cuda.is_available():
    for i in range(torch.cuda.device_count()):
        device_name = torch.cuda.get_device_properties(i).name
        cuda[f"CUDA:{i} {device_name}"] = f"cuda:{i}"



config_dir = "configs/"
def get_file_options(directory, extension):
    return [file for file in os.listdir(directory) if file.endswith(extension)]

def load_options():
    config_list = get_file_options(config_dir, ".json")
    return config_list

# ckpt_list, config_list, cluster_list, diff_list, diff_config_list = load_options()
config_list = load_options()


def load_json_encoder(config_choice, choice_ckpt):
    if config_choice == "no_config":
        return "未启用自动加载，请手动选择配置文件"
    if choice_ckpt == "no_model":
        return "请先选择模型"
    config_file = Config(os.path.join(config_dir, config_choice), "json")
    config = config_file.read()
    print(config)
    try:
        #比对配置文件中的模型维度与该encoder的实际维度是否对应，防止古神语
        config_encoder = config["model"].get("speech_encoder", "no_encoder")
        config_dim = config["model"]["ssl_dim"]
        #旧版配置文件自动匹配
        # if config_encoder == "no_encoder":
        #     config_encoder = config["model"]["speech_encoder"] = "vec256l9" if config_dim == 256 else "vec768l12"
        #     config_file.save(config)
        # correct_dim = MODEL_TYPE.get(config_encoder, "unknown")
        # if config_dim != correct_dim:
        #     return "配置文件中的编码器与模型维度不匹配"
        return config_encoder
    except Exception as e:
        return f"出错了: {e}"
    

app = gr.Blocks()
with app:
    gr.Markdown(value="""
        ### vits WebUI 推理&训练
        用于血腥、暴力、性相关、政治相关内容，会承担法律责任，这不是玩笑。
        """)
    with gr.Tabs():
        with gr.TabItem("训练"):
            gr.Markdown(value="""请将数据集文件夹放置在dataset_raw文件夹下，确认放置正确后点击下方获取数据集名称""")
            with gr.Row():
                config_choice = gr.Dropdown(label="配置文件", choices=config_list )
                config_info = gr.Textbox(label="配置文件编码器", placeholder="请选择配置文件")
            gr.Markdown(value="""**请检查模型和配置文件的编码器是否匹配**""")

            # config_choice.change(load_json_encoder, [config_choice], [config_info])

            with gr.Row():
                gr.Markdown(value="""**vec256l9**: ContentVec(256Layer9)，旧版本叫v1，So-VITS-SVC 4.0的基础版本，**不推荐使用**
                                **vec768l12**: 特征输入更换为ContentVec的第12层Transformer输出，模型理论上会更加还原训练集音色
                                **hubertsoft**: So-VITS-SVC 3.0使用的编码器，咬字更为准确，但可能存在多说话人音色泄露问题
                                **whisper-ppg**: 来自OpenAI，咬字最为准确，但和Hubertsoft一样存在多说话人音色泄露，且显存占用和训练时间有明显增加。
                                解锁更多编码器选项，请见[这里](https://www.yuque.com/umoubuton/ueupp5/kmui02dszo5zrqkz)
                """)
                gr.Markdown(value="""**crepe**: 抗噪能力最强，但预处理速度慢（不过如果你的显卡很强的话速度会很快）
                                **pm**: 预处理速度快，但抗噪能力较弱
                                **dio**: 先前版本预处理默认使用的f0预测器，比较拉胯不推荐使用
                                **harvest**: 有一定抗噪能力，预处理显存占用友好，速度比较慢
                                **rmvpe**: 最精准的预测器，crepe的完全上位替代
                """)
            with gr.Row():
                branch_selection = gr.Dropdown(label="选择训练使用的编码器", choices=[], value="vec768l12", interactive=True)
                f0_predictor_selection = gr.Dropdown(label="选择训练使用的f0预测器", choices=[], value="crepe", interactive=True)
            with gr.Row():
                use_diff = gr.Checkbox(label="是否使用浅扩散模型，如要训练浅扩散模型请勾选此项", value=True)
                vol_aug=gr.Checkbox(label="是否启用响度嵌入和音量增强，启用后可以根据输入源控制输出响度，但对数据集质量的要求更高。**仅支持vec768l12编码器**", value=False)
            with gr.Row():
                raw_preprocess=gr.Button("数据预处理", variant="primary")
                regenerate_config_btn=gr.Button("重新生成配置文件", variant="primary")
            preprocess_output=gr.Textbox(label="预处理输出信息，完成后请检查一下是否有报错信息，如无则可以进行下一步", max_lines=999)
            clear_preprocess_output=gr.Button("清空输出信息")
            with gr.Group():
                gr.Markdown(value="""填写训练设置和超参数""")
                with gr.Row():
                    gr.Textbox(label="当前使用显卡信息", value=gpu_info)
                    gpu_selection=gr.Textbox(label="多卡用户请指定希望训练使用的显卡ID（0,1,2...）", value=gpus, interactive=True)
                with gr.Row():
                    #tiny_enable = gr.Checkbox(label="是否启用TINY训练，TINY模型的显存占用更低，推理速度更快，但质量有所削减。仅支持vec768的响度嵌入", value=False)
                    gr.Markdown("请检查右侧的说话人列表是否和你要训练的目标说话人一致，确认无误后点击写入配置文件，然后就可以开始训练了")
                    speakers=gr.Textbox(label="说话人列表")
            with gr.Row():
                save_params=gr.Button("将当前设置保存为默认设置", variant="primary")
                write_config=gr.Button("写入配置文件", variant="primary")
            write_config_output=gr.Textbox(label="输出信息")

            gr.Markdown(value="""**点击从头开始训练**将会自动将已有的训练进度保存到models_backup文件夹，并自动装载预训练模型。
                **继续上一次的训练进度**将从上一个保存模型的进度继续训练。继续训练进度无需重新预处理和写入配置文件。
                关于扩散、聚类和特征检索的详细说明请看[此处](https://www.yuque.com/umoubuton/ueupp5/kmui02dszo5zrqkz)。
                """)
            with gr.Row():
                with gr.Column():
                    start_training=gr.Button("从头开始训练", variant="primary")
                    training_output=gr.Textbox(label="训练输出信息")
                with gr.Column():
                    continue_training_btn=gr.Button("继续上一次的训练进度", variant="primary")
                    continue_training_output=gr.Textbox(label="训练输出信息")
            with gr.Row():
                with gr.Column():
                    diff_training_btn=gr.Button("从头训练扩散模型", variant="primary")
                    diff_training_output=gr.Textbox(label="训练输出信息")
                with gr.Column():
                    diff_continue_training_btn=gr.Button("继续训练扩散模型", variant="primary")
                    diff_continue_training_output=gr.Textbox(label="训练输出信息") 
            with gr.Accordion(label = "聚类、特征检索训练", open=False):
                with gr.Row():               
                    with gr.Column():
                        kmeans_button=gr.Button("训练聚类模型", variant="primary")
                        kmeans_gpu = gr.Checkbox(label="使用GPU训练", value=True)
                        kmeans_output=gr.Textbox(label="训练输出信息")
                    with gr.Column():
                        index_button=gr.Button("训练特征检索模型", variant="primary")
                        index_output=gr.Textbox(label="训练输出信息")

        with gr.TabItem("小工具/实验室特性"):
            gr.Markdown(value="""
                        ### So-vits-svc 4.1 小工具/实验室特性
                        提供了一些有趣或实用的小工具，可以自行探索
                        """)
            
            with gr.Tabs():
                with gr.TabItem("静态声线融合"):
                    gr.Markdown(value="""
                        <font size=2> 介绍:该功能可以将多个声音模型合成为一个声音模型(多个模型参数的凸组合或线性组合)，从而制造出现实中不存在的声线 
                                          注意：
                                          1.该功能仅支持单说话人的模型
                                          2.如果强行使用多说话人模型，需要保证多个模型的说话人数量相同，这样可以混合同一个SpaekerID下的声音
                                          3.保证所有待混合模型的config.json中的model字段是相同的
                                          4.输出的混合模型可以使用待合成模型的任意一个config.json，但聚类模型将不能使用
                                          5.批量上传模型的时候最好把模型放到一个文件夹选中后一起上传
                                          6.混合比例调整建议大小在0-100之间，也可以调为其他数字，但在线性组合模式下会出现未知的效果
                                          7.混合完毕后，文件将会保存在项目根目录中，文件名为output.pth
                                          8.凸组合模式会将混合比例执行Softmax使混合比例相加为1，而线性组合模式不会
                        </font>
                        """)
                    mix_model_path = gr.Files(label="选择需要混合模型文件")
                    mix_model_upload_button = gr.UploadButton("选择/追加需要混合模型文件", file_count="multiple")
                    mix_model_output1 = gr.Textbox(
                                            label="混合比例调整，单位/%",
                                            interactive = True
                                         )
                    mix_mode = gr.Radio(choices=["凸组合", "线性组合"], label="融合模式",value="凸组合",interactive = True)
                    mix_submit = gr.Button("声线融合启动", variant="primary")
                    mix_model_output2 = gr.Textbox(
                                            label="Output Message"
                                         )

                with gr.TabItem("智能音频切片"):
                    gr.Markdown(value="""
                        该工具可以实现对音频的切片，无需调整参数即可完成符合要求的数据集制作。
                        数据集要求的音频切片约在2-15秒内，用传统的Slicer-GUI切片工具需要精准调参和二次切片才能符合要求，该工具省去了上述繁琐的操作，只要上传原始音频即可一键制作数据集。
                    """)
                    with gr.Row():
                        raw_audio_path = gr.Textbox(label="原始音频文件夹", placeholder="包含所有待切片音频的文件夹，示例: D:\干声\speakers")
                        load_raw_audio_btn = gr.Button("加载原始音频", variant = "primary")
                    load_raw_audio_output = gr.Textbox(label = "输出信息")
                    raw_audio_dataset = gr.Textbox(label = "音频列表", value = "")
                    slicer_output_dir = gr.Textbox(label = "输出目录", placeholder = "选择输出目录（不要和输入音频是同一个文件夹）")
                    with gr.Row():
                        process_method = gr.Radio(label = "对过短音频的处理方式", choices = ["丢弃","将过短音频整合为长音频"], value = "丢弃")
                        max_sec = gr.Number(label = "切片的最长秒数", value = 15)
                        min_sec = gr.Number(label = "切片的最短秒数", value = 2)
                    slicer_btn = gr.Button("开始切片", variant = "primary")
                    slicer_output_msg = gr.Textbox(label = "输出信息")

    app.queue(concurrency_count=1022, max_size=2044).launch(share=True, server_port=8080)



class Config:
    def __init__(self, path, type):
        self.path = path
        self.type = type
    
    def read(self):
        if self.type == "json":
            with open(self.path, 'r') as f:
                return json.load(f)
        if self.type == "yaml":
            with open(self.path, 'r') as f:
                return yaml.safe_load(f)
    
    def save(self, content):
        if self.type == "json":
            with open(self.path, 'w') as f:
                json.dump(content, f, indent=4)
        if self.type == "yaml":
            with open(self.path, 'w') as f:
                yaml.safe_dump(content, f, default_flow_style=False, sort_keys=False)